# backtesting/ma_strategy.py
import matplotlib.pyplot as plt
import pandas as pd

from data.mysql_db import DataManager
from indicators.indicator_calculator import TechnicalIndicators


class MovingAverageStrategy:

    def __init__(self, indicator: TechnicalIndicators, short_window=5, long_window=20):
        """
        初始化均线策略

        参数:

            short_window (int): 短期均线周期
            long_window (int): 长期均线周期
        """
        self.short_window = short_window
        self.long_window = long_window
        self.indicator = indicator

    def generate_signals(self) -> pd.DataFrame:
        """
        生成基于 MACD 的交易信号
        金叉（买入信号）：当 DIF（MACD 线）上穿 MACD（信号线）时，生成买入信号（Signal = 1）。
        死叉（卖出信号）：当 DIF 下穿 MACD 时，生成卖出信号（Signal = -1）。
        信号变化点：通过 Position = Signal.diff() 标记实际交易点位（Position = 1 表示买入，-1 表示卖出）。

        返回:
            DataFrame: 包含交易信号的 DataFrame
            返回的 DataFrame 包含以下关键列：
                trade_date: 交易日期
                close: 收盘价
                DIF: MACD 线
                MACD: 信号线
                Signal: 交易信号（1 买入，-1 卖出，0 无信号）
                Position: 信号变化点（用于触发交易）
        """

        # 计算 MACD 指标
        self.indicator.calculate_macd(
            short=self.short_window,
            long=self.long_window,
            signal=9  # 信号线周期（默认值）
        )

        df = self.indicator.data

        # 确保数据完整性（处理 NaN 值）
        df.dropna(inplace=True)

        # 生成信号
        df['Signal'] = 0  # 0 表示无信号

        # 金叉：MACD 上穿信号线
        df.loc[(df['DIF'].shift(1) < df['MACD'].shift(1)) &
               (df['DIF'] > df['MACD']), 'Signal'] = 1  # 买入信号

        # 死叉：MACD 下穿信号线
        df.loc[(df['DIF'].shift(1) > df['MACD'].shift(1)) &
               (df['DIF'] < df['MACD']), 'Signal'] = -1  # 卖出信号

        # 信号变化点（用于交易执行）
        df['Position'] = df['Signal'].diff()

        # 返回包含信号的 DataFrame
        return df[['trade_date', 'close', 'DIF', 'MACD', 'Signal', 'Position']]


class MACDBarStrategy:

    def __init__(self, indicator: TechnicalIndicators, short_window=12, long_window=26):
        """
        初始化 MACD 柱策略

        参数:
            short_window (int): 短期均线周期
            long_window (int): 长期均线周期
            indicator (TechnicalIndicators): 技术指标计算器
        """
        self.short_window = short_window
        self.long_window = long_window
        self.indicator = indicator

    def generate_signals(self) -> pd.DataFrame:
        """
        根据 MACD 柱的变化生成交易信号

        信号规则：
        1. DIF > 0 时才交易
        2. 买点：MACD 柱（Histogram）由负变正
        3. 卖点：MACD 柱由正变负

        返回:
            DataFrame: 包含交易信号的 DataFrame
        """
        # 计算 MACD 指标
        self.indicator.calculate_macd(
            short=self.short_window,
            long=self.long_window,
            signal=9  # 信号线周期（默认值）
        )

        df = self.indicator.data
        df.dropna(inplace=True)  # 处理 NaN 值

        # 计算 MACD 柱（Histogram）
        df['Histogram'] = df['DIF'] - df['MACD']

        # 初始化信号列
        df['Signal'] = 0  # 0 表示无信号

        # 买入条件：DIF > 0 且 Histogram 由负变正
        df.loc[
            (df['DIF'] > 0) &
            (df['Histogram'].shift(1) < 0) &
            (df['Histogram'] > 0),
            'Signal'
        ] = 1  # 买入信号

        # 卖出条件：DIF > 0 且 Histogram 由正变负
        df.loc[
            (df['DIF'] > 0) &
            (df['Histogram'].shift(1) > 0) &
            (df['Histogram'] < 0),
            'Signal'
        ] = -1  # 卖出信号

        # 信号变化点（用于交易执行）
        df['Position'] = df['Signal'].diff()

        # 返回结果
        return df[['trade_date', 'close', 'DIF', 'MACD', 'Histogram', 'Signal', 'Position']]


def plot_signals(df) -> None:
    """
    绘制买卖点
    :param df:
    :return:
    """
    plt.figure(figsize=(14, 7))
    plt.plot(df['close'], label='Close Price')
    plt.plot(df['DIF'], label='DIF')
    plt.plot(df['MACD'], label='MACD')

    # 标记买入信号（Histogram 由负变正）
    buy_mask = (df['Signal'] == 1) & (df['DIF'] > 0)
    plt.scatter(
        df[buy_mask].index,
        df['close'][buy_mask],
        color='red',
        label='Buy Signal',
        marker='^',
        alpha=1
    )

    # 标记卖出信号（Histogram 由正变负）
    sell_mask = (df['Signal'] == -1) & (df['DIF'] > 0)
    plt.scatter(
        df[sell_mask].index,
        df['close'][sell_mask],
        color='green',
        label='Sell Signal',
        marker='v',
        alpha=1
    )

    plt.legend()
    plt.show()


# 示例用法
if __name__ == '__main__':
    data_manager = DataManager()
    ts_code = '600610.SH'
    df = data_manager.get_stock_daily_data(ts_code)
    strategy = MovingAverageStrategy(TechnicalIndicators(df), short_window=5, long_window=20)
    signals = strategy.generate_signals()
    plot_signals(strategy.indicator.data)
